import pandas as pd
from sklearn.model_selection import train_test_split
import torch
from nlpx.tokenize.utils import get_df_text_labels
from nlpx.dataset import TextDataset, text_collate
from transformers_model import AutoCNNTextClassifier, ErnieRNNAttentionTextClassifier
from nlpx.model.wrapper import ClassifyModelWrapper

pretrained_path = "albert_chinese_tiny"
# pretrained_path = "nghuyong/ernie-3.0-base-zh"
file = "~/project/python/parttime/归档/text_gcn/data/北方地区不安全事件统计20240331.csv"

if __name__ == '__main__':
    df = pd.read_csv(file, encoding="GBK")
    texts, labels, classes = get_df_text_labels(
        df, text_col="故障描述", label_col="故障标志"
    )

    train_texts, test_texts, y_train, y_test = train_test_split(texts, labels, test_size=0.2, random_state=42)
    train_set = TextDataset(train_texts, y_train)
    test_set = TextDataset(test_texts, y_test)

    model = AutoCNNTextClassifier(pretrained_path, len(classes), num_train_layers=2)
    wrapper = ClassifyModelWrapper(model, classes)
    _ = wrapper.train(train_set, test_set, collate_fn=text_collate)
    